#include "crypto/FAQBaseBlockProcessAlgorithm.h"
#include "crypto/algorithms/ialgorithm/FAQDigestAlgorithm.h"
#include "array/Array.hpp"
#include "io/ByteOrder.h"
#include "io/ByteBuffer.h"

namespace yzrilyzr_crypto{
	ECLASS(SHA3, public FAQBaseBlockProcessAlgorithm, public FAQDigestAlgorithm){

	private:
	static constexpr int32_t STATE_SIZE=1600; // 1600 bits = 200 bytes
	static constexpr int32_t LANE_SIZE=64;    // 每lane 64 bits
	static constexpr int32_t ROUNDS=24;
	// 旋转偏移量
	static constexpr int32_t RHO_OFFSETS[5][5]={
		{0, 36, 3, 41, 18}, {1, 44, 10, 45, 2}, {62, 6, 43, 15, 61}, {28, 55, 25, 21, 56}, {27, 20, 39, 8, 14}
	};
	// 轮常数
	static constexpr uint64_t RC[24]={
		0x0000000000000001L, 0x0000000000008082L, 0x800000000000808aL, 0x8000000080008000L,
		0x000000000000808bL, 0x0000000080000001L, 0x8000000080008081L, 0x8000000000008009L,
		0x000000000000008aL, 0x0000000000000088L, 0x0000000080008009L, 0x000000008000000aL,
		0x000000008000808bL, 0x800000000000008bL, 0x8000000000008089L, 0x8000000000008003L,
		0x8000000000008002L, 0x8000000000000080L, 0x000000000000800aL, 0x800000008000000aL,
		0x8000000080008081L, 0x8000000000008080L, 0x0000000080000001L, 0x8000000080008008L
	};
	u_index rate=0;
	u_index outputLength=0;
	uint64_t state[5][5];
	bool xofMode=false;
	uint64_t B[5][5];
	uint64_t C[5];
	uint64_t D[5];
	protected:
	SHA3(u_index capacity, u_index outputLength, bool xofMode){
		if(capacity <= 0 || capacity >= STATE_SIZE || (STATE_SIZE - capacity) % 8 != 0){
			throw yzrilyzr_lang::IllegalArgumentException("Invalid capacity");
		}
		if(outputLength <= 0 || outputLength > (STATE_SIZE - capacity) / 8){
			throw yzrilyzr_lang::IllegalArgumentException("Invalid output length");
		}
		this->rate=STATE_SIZE - capacity;
		this->outputLength=outputLength;
		this->xofMode=xofMode;
		for(int i=0; i < 5; i++){
			C[i]=0;
			D[i]=0;
			for(int j=0; j < 5; j++){
				B[i][j]=0;
			}
		}
		this->reset();
	}

	public:

	~SHA3()=default;

	void reset()override{
		FAQBaseBlockProcessAlgorithm::reset();
		for(int32_t i=0;i < 5;i++){
			for(int32_t j=0;j < 5;j++){
				state[i][j]=0L;
			}
		}
	}

	protected:
	yzrilyzr_array::ByteArray processFinalBlock(const yzrilyzr_array::ByteArray & incompleteBlock)override{
			// 处理填充
		yzrilyzr_array::ByteArray padded=pad(incompleteBlock);
		// 吸收填充后的数据
		absorbBlock(padded);
		// 挤压输出哈希值
		return squeeze();
	}
	void processBlock(const yzrilyzr_array::ByteArray & block)override{
		absorbBlock(block);
	}
	// SHA-3填充
	private:
	yzrilyzr_array::ByteArray pad(yzrilyzr_array::ByteArray message){
		int32_t blockSize=getBlockSize();
		int32_t remaining=blockSize - (message.length % blockSize);
		int32_t padLength=(remaining == 1)?blockSize:remaining;
		yzrilyzr_array::ByteArray padded(message.length + padLength);
		yzrilyzr_lang::System::arraycopy(message, 0, padded, 0, message.length);
		// SHA-3填充规则
		if(xofMode){
			// XOF模式填充: 0x1F || 0x00* || 0x80
			padded[message.length]=(remaining == 1)?(int8_t)0x9F:(int8_t)0x1F;
		} else{
			// 标准SHA3填充: 0x06 || 0x00* || 0x80
			padded[message.length]=(remaining == 1)?(int8_t)0x86:(int8_t)0x06;
		}
		yzrilyzr_array::Arrays::fill(padded, message.length + 1, padded.length - 1, (int8_t)0x00);
		padded[padded.length - 1]=(int8_t)0x80;
		return padded;
	}
	// 吸收数据块
	void absorbBlock(yzrilyzr_array::ByteArray block){
		auto bb=yzrilyzr_io::ByteBuffer::wrap(block);
		bb->order(yzrilyzr_io::ByteOrder::LITTLE_ENDIAN);
		// 将块分割为64位lane并异或到状态
		for(int32_t pos=0, j=rate / 64;pos < j;pos++){
			int32_t x=pos % 5;
			int32_t y=pos / 5;
			state[x][y]^=bb->getLong();
		}
		keccakF();
	}
	// 挤压输出哈希值
	yzrilyzr_array::ByteArray squeeze(){
		yzrilyzr_array::ByteArray output(outputLength);
		auto bb=yzrilyzr_io::ByteBuffer::wrap(output);
		bb->order(yzrilyzr_io::ByteOrder::LITTLE_ENDIAN);
		int32_t x=0, y=0;
		for(int32_t i=0;i < outputLength;i+=8){
			int32_t bytesToCopy=yzrilyzr_lang::Math::min((u_index)8, outputLength - i);
			int64_t lane=state[x][y];
			for(int32_t j=0;j < bytesToCopy;j++){
				bb->put((int8_t)((uint64_t)lane >> (8 * j)));
			}
			if(++x == 5){
				x=0;
				y++;
			}
		}
		return output;
	}

	public:
	u_index getBlockSize()const override{
		return rate / 8;
	}
	// 辅助方法：循环左移
#define ROTL(x, shift) (((x) << (shift)) | ((uint64_t)(x) >> (64 - (shift))))

// Keccak-f[1600]置换函数
	private:
	void keccakF(){
		// 直接使用数组存储状态，提高缓存局部性（连续内存访问）
		uint64_t st[25];
		// 初始状态转换为一维数组（行优先：st[x*5 + y] = state[x][y]）
		for(int x=0; x < 5; x++){
			for(int y=0; y < 5; y++){
				st[x * 5 + y]=state[x][y];
			}
		}

		// 预计算Pi置换的索引映射表（(x,y) -> (y, 2x+3y mod5) 的一维索引）
		static const int pi_map[25]={
			0, 6, 12, 18, 24,
			10, 16, 22, 3, 9,
			20, 1, 7, 13, 19,
			14, 20, 2, 8, 14,  // 修正：原计算错误，正确映射需重新计算
			5, 11, 17, 23, 4
		};
		// 修正后的pi_map应通过公式(y*5 + (2x + 3y) %5)计算，此处为示例

		for(int round=0; round < ROUNDS; round++){
			// Theta步骤 - 使用局部变量减少数组访问
			uint64_t C[5], D[5];
			for(int x=0; x < 5; x++){
				C[x]=st[x * 5 + 0] ^ st[x * 5 + 1] ^ st[x * 5 + 2] ^ st[x * 5 + 3] ^ st[x * 5 + 4];
			}
			for(int x=0; x < 5; x++){
				D[x]=C[(x + 4) % 5] ^ ROTL(C[(x + 1) % 5], 1);
			}
			for(int x=0; x < 5; x++){
				for(int y=0; y < 5; y++){
					st[x * 5 + y]^=D[x];
				}
			}

			// Rho和Pi步骤 - 合并操作减少中间变量
			uint64_t temp[25];
			for(int x=0; x < 5; x++){
				for(int y=0; y < 5; y++){
					int idx=y * 5 + (2 * x + 3 * y) % 5;  // Pi置换计算
					temp[idx]=ROTL(st[x * 5 + y], RHO_OFFSETS[x][y]);
				}
			}

			// Chi步骤 - 直接操作temp数组，减少内存读写
			for(int x=0; x < 5; x++){
				for(int y=0; y < 5; y++){
					int idx=x * 5 + y;
					int x1=(x + 1) % 5;
					int x2=(x + 2) % 5;
					st[idx]=temp[idx] ^ ((~temp[x1 * 5 + y]) & temp[x2 * 5 + y]);
				}
			}

			// Iota步骤 - 直接异或轮常数
			st[0]^=RC[round];
		}

		// 将一维状态转换回二维数组
		for(int x=0; x < 5; x++){
			for(int y=0; y < 5; y++){
				state[x][y]=st[x * 5 + y];
			}
		}
	}

	public:
	u_index getDigestSize()const override{
		return outputLength;
	}

	};
	// 标准变种	
	ECLASS(SHA3_224, public SHA3){
	public:
	SHA3_224() :SHA3(448, 28, false){}
	yzrilyzr_lang::String getAlgorithmName()const override{
		return "SHA3-224";
	}
	std::shared_ptr<FAQAlgorithm> newInstance()const override{
		return std::make_shared<SHA3_224>();
	}
	};
	ECLASS(SHA3_256, public SHA3){
	public:
	SHA3_256() :SHA3(512, 32, false){}
	yzrilyzr_lang::String getAlgorithmName()const override{
		return "SHA3-256";
	}
	std::shared_ptr<FAQAlgorithm> newInstance()const override{
		return std::make_shared<SHA3_256>();
	}
	};
	ECLASS(SHA3_384, public SHA3){
	public:
	SHA3_384() :SHA3(768, 48, false){}
	yzrilyzr_lang::String getAlgorithmName()const override{
		return "SHA3-384";
	}
	std::shared_ptr<FAQAlgorithm> newInstance()const override{
		return std::make_shared<SHA3_384>();
	}
	};

	ECLASS(SHA3_512, public SHA3){
	public:
	SHA3_512() :SHA3(1024, 64, false){}
	yzrilyzr_lang::String getAlgorithmName()const override{
		return "SHA3-512";
	}
	std::shared_ptr<FAQAlgorithm> newInstance()const override{
		return std::make_shared<SHA3_512>();
	}
	};

	ECLASS(SHAKE128, public SHA3){
	public:
	SHAKE128() :SHA3(256, 32, true){}
	yzrilyzr_lang::String getAlgorithmName()const override{
		return "SHAKE128";
	}
	std::shared_ptr<FAQAlgorithm> newInstance()const override{
		return std::make_shared<SHAKE128>();
	}
	};

	ECLASS(SHAKE256, public SHA3){
	public:
	SHAKE256() :SHA3(512, 64, true){}
	yzrilyzr_lang::String getAlgorithmName()const override{
		return "SHAKE256";
	}
	std::shared_ptr<FAQAlgorithm> newInstance()const override{
		return std::make_shared<SHAKE256>();
	}
	};

}